"""
Provide basic validation operations. Which means, consume
traversals we get.
"""

from ast import *
import ast  # for AST rebuild
import json
import astor
from core.query_set import QuerySet
from core.utility import vertex_attr, rebuild_ast, vertex_attr_str
from core.dfg_graph import DFGNode, DFGPrinter, DFGVisitor, build_dfg_graph
from .traverse import q_dfg_paths


def _get_dict_key_set(value: Dict):
    key_set = []
    for key in value.keys:
        if isinstance(key, Constant):
            key_set.append(key.value)
    return key_set


class TorchSaveVisitor(DFGVisitor):
    def __init__(self):
        self.status = True
        self.key_set = []
        self.bad_node = None

    def enter(self, node: DFGNode):
        return self.visit(node)

    def exit(self, node: DFGNode):
        return self.status

    def visit(self, node: DFGNode):
        vertex = node.vertex
        if vertex is None:
            return True
        code = vertex_attr_str(vertex, "code")
        ast_node = rebuild_ast(vertex_attr(vertex, "ast"))
        if code.startswith("torch.save("):
            try:
                arg = ast_node.value.args[0]
                if isinstance(arg, Dict):
                    self.key_set = _get_dict_key_set(arg)
                    self.bad_node = node
                    self.status = False
            except:
                pass
        elif isinstance(ast_node, Assign) and isinstance(ast_node.value, Dict):
            self.key_set = _get_dict_key_set(ast_node.value)
            # we stop at the first assignment with Dict
            self.bad_node = node
            self.status = False
        return self.status


def _vertex_back_trace(node):
    vertices = []
    if node is not None:
        while node.parent is not None:
            vertices.append(node.vertex)
            node = node.parent
    for vertex in reversed(vertices):
        yield vertex


def c_torch_save_consumer(q: QuerySet):
    # in case some unwanted changes happen
    g = q.query.clone()

    root = build_dfg_graph(g.to_list())
    
    # root.accept(DFGPrinter())
    visitor = TorchSaveVisitor()
    root.accept(visitor)
    # check if all expected keys are present
    EXPECTED_KEYS = ["epoch", "state_dict", "optimizer", "loss"]
    missing_keys = []
    used_keys = []
    for key in EXPECTED_KEYS:
        flag = False
        for actual_key in visitor.key_set:
            if actual_key in used_keys:
                continue
            if actual_key.find(key) != -1:
                used_keys.append(actual_key)
                flag = True
                break
        if not flag:
            missing_keys.append(key)

    if len(missing_keys) > 0:
        print("torch.save() missing keys: " + str(missing_keys))
        for vertex in _vertex_back_trace(visitor.bad_node):
            print("    Line " + vertex_attr_str(vertex, "lineno") + ":")
            print("        " + vertex_attr_str(vertex, "code"))
        return True
    else:
        print("torch.save() is good")
        return False
